source("./function_code/relabel_vars.R")

### Function for coefficients
msmCoef = function(themodel) {
  # Extract coefficients
  vs = names(hazard.msm(themodel))
  
  hrList = list()
  for (i in 1:length(vs)){
    varRes = hazard.msm(themodel)[vs[i]]
    
    # Get the state transition labels
    stl = names(varRes[[1]][,1])
    
    # Get hazard ratios and 95% limits
    hr = log(varRes[[1]][,1])
    hrl = log(varRes[[1]][,2])
    hru = log(varRes[[1]][,3])
    
    hrs = data.frame(vars=vs[i], trans=stl, est=hr, lower=hrl, upper=hru)
    hrList[[i]] = hrs
  }
  hrest = do.call("rbind", hrList)
  
  return(hrest)
}

# Generate coefficient plot 
msmCoefPlot = function(theresults, xmin, xmax, fileID) {
  # Make the transition labels a factor for proper ordering
  theresults[,"trans"] = factor(theresults[,"trans"], 
                                levels=c("State 1 - State 2",
                                         "State 2 - State 1",
                                         "State 1 - State 3",
                                         "State 2 - State 3"))
  
  # Add statistical significance variable
  theresults[,"stat"] = ifelse(theresults[,"lower"] > 0 | theresults[,"upper"] < 0, "Yes", "No")
  
  # Replace the variable names
  theresults = relabelVars(theresults)
  
  # Make the plot 
  theplot = ggplot(theresults, aes(y=fct_rev(vars), x=est, group=fct_rev(trans))) + 
    geom_vline(xintercept=0, linetype=2) + 
    geom_pointrange(aes(xmin=lower, xmax=upper, color=trans, shape=stat), position=position_dodge(width=0.6)) + 
    theme_bw() + xlab("Coefficient estimate") + ylab("") + 
    scale_color_manual("Transition", labels=c("Start neg.", "End neg.", "Victory/defeat", "Neg. settlement"), 
                       values=c("gray60", "gray40", "gray20", "black")) +
    scale_shape_manual("Statistically\nsignificant (95%)", values=c(16, 15)) +
    theme(legend.position="bottom", 
          legend.direction = "vertical") +
    guides(color = guide_legend(order=1),
           shape = guide_legend(order=2))
  
  print(theplot)
  
  # Make file name
  fn = paste0("./figures/msmCoefPlot", fileID, ".pdf")
  
  # Create file
  ggsave(fn, width=4.5, height=7)
}



# Generate coefficient plot for disaggregated negotiations
msmCoefPlot2 = function(theresults, xmin, xmax, fileID) {
  # Make the transition labels a factor for proper ordering
  theresults[,"trans"] = factor(theresults[,"trans"], 
                                levels=c("State 1 - State 2",
                                         "State 2 - State 1",
                                         "State 1 - State 3",
                                         "State 3 - State 1"))
  
  # Remove results for talks ending
  theresults = theresults %>% filter(trans %in% c("State 1 - State 2", "State 1 - State 3"))
  
  
  # Add statistical significance variable
  theresults[,"stat"] = ifelse(theresults[,"lower"] > 0 | theresults[,"upper"] < 0, "Yes", "No")
  
  # Replace the variable names
  theresults = relabelVars(theresults)
  
  # Make the plot 
  theplot = ggplot(theresults, aes(y=fct_rev(vars), x=est, group=fct_rev(trans))) + 
    geom_vline(xintercept=0, linetype=2) + 
    geom_pointrange(aes(xmin=lower, xmax=upper, color=trans, shape=stat), position=position_dodge(width=0.6)) + 
    theme_bw() + xlab("Coefficient estimate") + ylab("") + 
    scale_shape_manual("Statistically\nsignificant (95%)", values=c(16, 15)) + 
    scale_color_manual("Start negotiation", labels=c("Internal", "External"), 
                       values=c("gray10", "gray60")) +
    theme(legend.position="bottom", 
          legend.direction = "vertical")
  
 
  print(theplot)
  
  # Make file name
  fn = paste0("./figures/msmCoefPlot", fileID, ".pdf")
  
  # Create file
  ggsave(fn, width=4.5, height=7)
}
